Skip to content

feat: add aa.interp_2d (NumPy + JAX bilinear interpolation) - #308

Merged
Jammy2211 merged 2 commits into
mainfrom
feature/jax-interp-2d
May 14, 2026
Merged

Jammy2211 merged 2 commits into
mainfrom
feature/jax-interp-2d

Conversation

@Jammy2211

@Jammy2211 Jammy2211 commented May 14, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Add a JAX-compatible 2D regular-grid bilinear interpolator at autoarray/operators/interp_2d.py, with both NumPy (scipy RegularGridInterpolator) and JAX (jax.scipy.ndimage.map_coordinates) paths and a single dispatcher surfaced at the top level as aa.interp_2d. Step 4 of the ellipse_fitting_jax feature — DatasetInterp in PyAutoGalaxy currently blocks JAX on scipy.interpolate, and this helper unblocks it.

API Changes

Adds a new public function autoarray.interp_2d(points, x_axis, y_axis, values, fill_value=0.0, xp=np) (also reachable as autoarray.operators.interp_2d.interp_2d). NumPy path matches scipy.interpolate.RegularGridInterpolator(bounds_error=False, fill_value=0.0) semantics byte-for-byte; JAX path uses jax.scipy.ndimage.map_coordinates(order=1, mode='constant') and is parity-tested against the NumPy path to rtol=1e-6 on random in-bounds queries. No symbols removed, renamed, or changed signature. See full details below.

Test Plan

  • python -m pytest test_autoarray/operators/test_interp_2d.py -v — 4/4 pass
  • Existing test_autoarray/ suite continues to pass
Full API Changes (for automation & release notes)

Added

  • autoarray.interp_2d(points, x_axis, y_axis, values, fill_value=0.0, xp=np) — 2D regular-grid bilinear interpolation. points is (N, 2) with column 0 → x_axis, column 1 → y_axis (matching scipy RegularGridInterpolator(points=(x_axis, y_axis)) convention). xp=np selects scipy; xp=jnp selects jax.scipy.ndimage.map_coordinates. Out-of-bounds queries return fill_value. Module: autoarray.operators.interp_2d.

Removed

None.

Changed Signature

None.

Changed Behaviour

None.

Migration

N/A — purely additive.

Generated with Claude Code

Adds autoarray/numerics/ subpackage with interp_2d.py: a 2D regular-grid
bilinear interpolator with NumPy (scipy RegularGridInterpolator) and JAX
(jax.scipy.ndimage.map_coordinates, order=1) backends behind a single
dispatcher interp_2d(points, x_axis, y_axis, values, fill_value=0.0, xp=np).

Unblocks DatasetInterp in PyAutoGalaxy for the ellipse_fitting_jax feature.

Issue PyAutoArray#306.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Downstream PR: PyAutoLabs/PyAutoGalaxy#398

Relocates interp_2d from the new autoarray/numerics/ subpackage into the
existing operators/ subpackage. operators/ already houses similar
grid-data operations (Convolver, contour, transformer, over_sampling),
so a dedicated numerics/ subpackage for a single function was
over-engineering. Top-level surface is now aa.interp_2d (mirroring how
aa.Convolver is exposed from operators/convolver.py).

Issue PyAutoArray#306.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211 Jammy2211 changed the title feat: add aa.numerics.interp_2d (NumPy + JAX bilinear interpolation) feat: add aa.interp_2d (NumPy + JAX bilinear interpolation) May 14, 2026
@Jammy2211
Jammy2211 merged commit 601f2af into main May 14, 2026
4 checks passed
@Jammy2211
Jammy2211 deleted the feature/jax-interp-2d branch May 14, 2026 11:25
@Jammy2211 Jammy2211 removed the pending-release PR queued for the next release build label Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant